// 沿 x 恒速 + z 力控（0–1N）；启动先去到指定 x,z 目标，x>=stop_x 后停止 test123
#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <eigen3/Eigen/Dense>
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

using namespace std::chrono_literals;
static const rclcpp::Logger LOGGER = rclcpp::get_logger("force_servo_control_node");

class ForceServoControlNode : public rclcpp::Node
{
public:
  ForceServoControlNode() : Node("force_servo_control_node")
  {
    // —— 主流程参数 ——
    vx_ = declare_parameter<double>("vx", 0.02);                   // x 向恒速 m/s
    stop_x_threshold_ = declare_parameter<double>("stop_x", 0.1);  // 到此停止 m
    publish_period_ms_ = declare_parameter<int>("period_ms", 40);  // 25 Hz
    cmd_topic_ = declare_parameter<std::string>("cmd_topic", "force_servo_control_node/delta_twist_cmds");

    // —— 力控（z）0–1 N ——（按需改）
    F_low_  = declare_parameter<double>("F_low", 0.0);
    F_high_ = declare_parameter<double>("F_high", 1.0);
    k_f_    = declare_parameter<double>("k_f", 0.005);      // m/(s·N)
    v_z_max_= declare_parameter<double>("v_z_max", 0.02);   // m/s
    down_sign_ = declare_parameter<double>("down_sign", -1.0); // 往下方向（+z向上则取-1）

    // —— 目标点（一次性）——
    go_to_target_ = declare_parameter<bool>("go_to_target", true);
    target_x_ = declare_parameter<double>("target_x", 0.05);  // m（base_link）
    target_z_ = declare_parameter<double>("target_z", 0.20);  // m
    pos_k_    = declare_parameter<double>("pos_k", 1.0);      // 位置 P 增益（m/s per m）
    pos_vmax_ = declare_parameter<double>("pos_vmax", 0.05);  // 目标阶段速度限幅 m/s
    pos_tol_  = declare_parameter<double>("pos_tol", 0.003);  // 到达判定阈值 m

    // 订阅力数据（只用 z）
    force_sub_ = create_subscription<geometry_msgs::msg::TwistStamped>(
        "/force_data", 10,
        std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

    // TF
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

    // 发布给 Servo 的 Twist
    twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>(cmd_topic_, 10);
  }

  void initialize()
  {
    setupPlanningSceneMonitor();
    setupServo();

    state_ = go_to_target_ ? State::GO_TO_TARGET : State::RUN;

    timer_ = create_wall_timer(
        std::chrono::milliseconds(publish_period_ms_),
        std::bind(&ForceServoControlNode::control_loop, this));

    RCLCPP_INFO(LOGGER, "启动完成：vx=%.3f, stop_x=%.3f, F=[%.1f,%.1f]N, k_f=%.4f, v_z_max=%.3f, down_sign=%.1f",
                vx_, stop_x_threshold_, F_low_, F_high_, k_f_, v_z_max_, down_sign_);
    RCLCPP_INFO(LOGGER, "go_to_target=%s, target_x=%.3f, target_z=%.3f, pos_k=%.3f, pos_vmax=%.3f, pos_tol=%.3f",
                go_to_target_ ? "true":"false", target_x_, target_z_, pos_k_, pos_vmax_, pos_tol_);
  }

private:
  // —— 通信对象 ——
  rclcpp::Subscription<geometry_msgs::msg::TwistStamped>::SharedPtr force_sub_;
  rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
  std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
  std::shared_ptr<planning_scene_monitor::PlanningSceneMonitor> planning_scene_monitor_;
  std::unique_ptr<moveit_servo::Servo> servo_;
  rclcpp::TimerBase::SharedPtr timer_;

  // —— 状态 & 变量 ——
  enum class State { GO_TO_TARGET, RUN, STOPPED };
  State state_{State::RUN};
  Eigen::Vector3d current_force_{Eigen::Vector3d::Zero()};

  // —— 参数 ——
  double vx_;
  double stop_x_threshold_;
  int publish_period_ms_;
  std::string cmd_topic_;

  // 力控
  double F_low_, F_high_, k_f_, v_z_max_, down_sign_;

  // 目标点
  bool   go_to_target_;
  double target_x_, target_z_, pos_k_, pos_vmax_, pos_tol_;

  // ———— 回调：力数据（只用 z）————
  void force_callback(const geometry_msgs::msg::TwistStamped::SharedPtr msg)
  {
    current_force_.x() = 0.0;
    current_force_.y() = 0.0;
    current_force_.z() = msg->twist.linear.z;  // 单位 N
  }

  // ———— 主循环 ————
  void control_loop()
  {
    // 读取当前末端在 base_link 下的位置
    double cur_x = 0.0, cur_z = 0.0;
    try {
      auto tf = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
      cur_x = tf.transform.translation.x;
      cur_z = tf.transform.translation.z;
    } catch (tf2::TransformException &ex) {
      RCLCPP_WARN(LOGGER, "获取末端位置失败：%s", ex.what());
      return;
    }

    auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
    msg->header.stamp = this->now();
    msg->header.frame_id = "base_link";

    if (state_ == State::GO_TO_TARGET)
    {
      // 只在 x、z 两个方向靠近目标（y 不动）
      double ex = target_x_ - cur_x;
      double ez = target_z_ - cur_z;

      double vx_cmd = clamp(pos_k_ * ex, -pos_vmax_, pos_vmax_);
      double vz_cmd = clamp(pos_k_ * ez, -pos_vmax_, pos_vmax_);

      // 到达判据（x、z 都进 tol）
      if (std::abs(ex) < pos_tol_ && std::abs(ez) < pos_tol_)
      {
        vx_cmd = 0.0; vz_cmd = 0.0;
        state_ = State::RUN;
        RCLCPP_INFO(LOGGER, "已到达目标 (x=%.3f, z=%.3f)，切换到 RUN", cur_x, cur_z);
      }

      msg->twist.linear.x = vx_cmd;
      msg->twist.linear.y = 0.0;
      msg->twist.linear.z = vz_cmd;
      msg->twist.angular.x = 0.0;
      msg->twist.angular.y = 0.0;
      msg->twist.angular.z = 0.0;

      twist_cmd_pub_->publish(std::move(msg));
      return;
    }

    if (state_ == State::RUN)
    {
      // 到达 x 阈值就全停
      if (cur_x >= stop_x_threshold_) {
        state_ = State::STOPPED;
        RCLCPP_INFO(LOGGER, "末端 x=%.3f >= %.3f，进入 STOPPED", cur_x, stop_x_threshold_);
      }
    }

    if (state_ == State::STOPPED)
    {
      // 全停
      msg->twist.linear.x = 0.0;
      msg->twist.linear.y = 0.0;
      msg->twist.linear.z = 0.0;
      msg->twist.angular.x = 0.0;
      msg->twist.angular.y = 0.0;
      msg->twist.angular.z = 0.0;
      twist_cmd_pub_->publish(std::move(msg));
      return;
    }

    // —— RUN 状态：x 恒速 + z 力控（0–1N 死区）——
    const double Fz = current_force_.z();
    double v_z_cmd = 0.0;

    if (Fz < F_low_) {
      // 力偏小 => 往下
      v_z_cmd = down_sign_ * k_f_ * (F_low_ - Fz);
    } else if (Fz > F_high_) {
      // 力偏大 => 往上
      v_z_cmd = down_sign_ * k_f_ * (F_high_ - Fz);  // 注意这里是负数
    }
    // 限幅
    if (v_z_cmd >  v_z_max_) v_z_cmd =  v_z_max_;
    if (v_z_cmd < -v_z_max_) v_z_cmd = -v_z_max_;

    msg->twist.linear.x = vx_;       // 沿 x 恒速
    msg->twist.linear.y = 0.0;
    msg->twist.linear.z = v_z_cmd;   // z 力控
    msg->twist.angular.x = 0.0;
    msg->twist.angular.y = 0.0;
    msg->twist.angular.z = 0.0;

    twist_cmd_pub_->publish(std::move(msg));
  }

  // —— 工具：限幅 ——
  static double clamp(double v, double lo, double hi)
  { return std::max(lo, std::min(v, hi)); }

  // —— PSM & Servo ——
  void setupPlanningSceneMonitor()
  {
    planning_scene_monitor_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
        shared_from_this(), "robot_description", tf_buffer_, "planning_scene_monitor");

    if (planning_scene_monitor_->getPlanningScene())
    {
      planning_scene_monitor_->startStateMonitor("/joint_states");
      planning_scene_monitor_->setPlanningScenePublishingFrequency(25);
      planning_scene_monitor_->startPublishingPlanningScene(
          planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
          "/moveit_servo/publish_planning_scene");
      planning_scene_monitor_->startSceneMonitor();
      planning_scene_monitor_->providePlanningSceneService();
    }
    else
    {
      RCLCPP_FATAL(LOGGER, "PlanningSceneMonitor 配置失败");
      throw std::runtime_error("Planning scene setup failed");
    }
  }

  void setupServo()
  {
    auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(shared_from_this());
    if (!servo_parameters)
    {
      RCLCPP_FATAL(LOGGER, "无法加载 MoveIt Servo 参数");
      throw std::runtime_error("Failed to load servo parameters");
    }

    servo_ = std::make_unique<moveit_servo::Servo>(shared_from_this(), servo_parameters, planning_scene_monitor_);
    servo_->start();
    RCLCPP_INFO(LOGGER, "MoveIt Servo 已启动，监听话题：%s", cmd_topic_.c_str());
  }
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ForceServoControlNode>();
  node->initialize();   // 构造后再初始化，避免 bad_weak_ptr
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}